import numpy as np
import cv2
import os
import imageio
import datetime
import math
import openpyxl
from lib import Heading
import pandas as pd

timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
out_dir = f"VO_Output_{timestamp}"
os.makedirs(out_dir, exist_ok=True)

# 画布配置
scale = 50
img_size = 600

# 初始化位置和参数
robot_pos = np.array([0.0, 0.0])
goal = np.array([10.0, 10.0])
robot_radius = 0.1
robot_speed = 1.5
dt = 0.1
arrived_threshold = 0.5
path = [robot_pos.copy()]

obstacles = [
    {"pos": np.array([5.0, 5.0]), "vel": np.array([0.1, 0]), "radius": 0.5},
    {"pos": np.array([6.0, 8.0]), "vel": np.array([-0.05, -0.05]), "radius": 0.5},
    {"pos": np.array([7.0, 5.0]), "vel": np.array([0, 0.1]), "radius": 0.5},
]

trajectory = []

def world_to_screen(pos):
    return tuple((pos * scale ).astype(int))

def draw_scene(img, robot_pos, goal, obstacles):
    img[:] = 255
    cv2.circle(img, world_to_screen(goal), 8, (0, 0, 255), -1)
    cv2.circle(img, world_to_screen(robot_pos), int(robot_radius * scale), (255, 0, 255), -1)
    for obs in obstacles:
        cv2.circle(img, world_to_screen(obs["pos"]), int(obs["radius"] * scale), (0, 0, 0), 2)
        cv2.circle(img, world_to_screen(obs["pos"]), int(0.1 * scale), (0, 0, 0), -1)

    for i in range(1, len(trajectory)):
        cv2.line(img, world_to_screen(trajectory[i-1]), world_to_screen(trajectory[i]), (255, 0, 0), 2)


def will_collide(robot_pos, v, obs_pos, obs_vel, robot_radius, obs_radius, dt=0.1, predict_time=1.2):
    for t in np.arange(0, predict_time, dt):
        future_robot = robot_pos + v * t
        future_obs = obs_pos + obs_vel * t
        dist = np.linalg.norm(future_robot - future_obs)
        if dist < (robot_radius + obs_radius):
            return True
    return False

def is_velocity_safe(robot_pos, v, obstacles, robot_radius):
    for obs in obstacles:
        if will_collide(robot_pos, v, obs["pos"], obs["vel"], robot_radius, obs["radius"]):
            return False
    return True

def select_velocity(robot_pos, goal, obstacles, robot_radius, speed):
    best_v = None
    min_cost = float("inf")
    for angle in np.linspace(0, 2*np.pi, 72):
        v = speed * np.array([np.cos(angle), np.sin(angle)])
        if is_velocity_safe(robot_pos, v, obstacles, robot_radius):
            to_goal = goal - robot_pos
            cost = np.linalg.norm(to_goal - v)
            if cost < min_cost:
                min_cost = cost
                best_v = v
    return best_v

# 图像帧列表
frames = []
img = np.ones((img_size, img_size, 3), dtype=np.uint8) * 255

vel_log = []
heading_records = []
speed_records = []
# 主循环
while np.linalg.norm(robot_pos - goal) > arrived_threshold:
    draw_scene(img, robot_pos, goal, obstacles)

    v = select_velocity(robot_pos, goal, obstacles, robot_radius, robot_speed)
    if v is None:
        print("⚠️ 没有安全速度可用，停止规划")
        break

    robot_pos += v * dt
    trajectory.append(robot_pos.copy())

    for obs in obstacles:
        obs["pos"] += obs["vel"]*dt

    # for pt in trajectory:
    #     cv2.circle(img, world_to_screen(pt), 1, (128, 0, 128), -1)

    cv2.imshow("VO Simulation", img)
    frames.append(img.copy())

    if cv2.waitKey(1) == 27 or np.linalg.norm(robot_pos - goal) < 0.3:
        print("✅ 到达目标或中止")
        break

    speed = np.linalg.norm(v)
    angle_deg = math.degrees(math.atan2(v[1], v[0]))
    vel_log.append([speed, angle_deg])
    heading_records.append(angle_deg)
    speed_records.append(speed)

cv2.destroyAllWindows()

# 保存GIF
gif_path = os.path.join(out_dir, "trajectory.gif")
imageio.mimsave(gif_path, frames, duration=0.05)

evaluator = Heading.PathSmoothnessEvaluator(heading_records, unit='deg')
result = evaluator.evaluate()
print(f"累计转角: {result['total_turning_angle']:.2f}°")
print(f"最大单次转角: {result['max_turning_angle']:.2f}°")

# 保存Excel
timestamps = [i for i in range(len(heading_records))]
df_records = pd.DataFrame({
    'Time': timestamps,
    'Speed': speed_records,
    'Yaw': heading_records

})
df_records.to_excel(os.path.join(out_dir, 'velocity_heading.xlsx'), index=False)


print(f"✅ 动图和航速航向数据已保存至文件夹：{out_dir}")